home *** CD-ROM | disk | FTP | other *** search
/ Personal Computer World 2008 February / PCWFEB08.iso / Software / Freeware / Miro 1.0 / Miro_Installer.exe / xulrunner / python / BitTorrent / Choker.py < prev    next >
Encoding:
Python Source  |  2007-11-12  |  9.7 KB  |  423 lines
  1. # Written by Bram Cohen
  2. # see LICENSE.txt for license information
  3.  
  4. from random import randrange
  5.  
  6. class Choker:
  7.     def __init__(self, max_uploads, schedule, done = lambda: False, min_uploads = None):
  8.         self.max_uploads = max_uploads
  9.         if min_uploads is None:
  10.             min_uploads = max_uploads
  11.         self.min_uploads = min_uploads
  12.         self.schedule = schedule
  13.         self.connections = []
  14.         self.count = 0
  15.         self.done = done
  16.         schedule(self._round_robin, 10)
  17.     
  18.     def _round_robin(self):
  19.         self.schedule(self._round_robin, 10)
  20.         self.count += 1
  21.         if self.count % 3 == 0:
  22.             for i in xrange(len(self.connections)):
  23.                 u = self.connections[i].get_upload()
  24.                 if u.is_choked() and u.is_interested():
  25.                     self.connections = self.connections[i:] + self.connections[:i]
  26.                     break
  27.         self._rechoke()
  28.  
  29.     def _snubbed(self, c):
  30.         if self.done():
  31.             return False
  32.         return c.get_download().is_snubbed()
  33.  
  34.     def _rate(self, c):
  35.         if self.done():
  36.             return c.get_upload().get_rate()
  37.         else:
  38.             return c.get_download().get_rate()
  39.  
  40.     def _rechoke(self):
  41.         preferred = []
  42.         for c in self.connections:
  43.             if not self._snubbed(c) and c.get_upload().is_interested():
  44.                 preferred.append((-self._rate(c), c))
  45.         preferred.sort()
  46.         del preferred[self.max_uploads - 1:]
  47.         preferred = [x[1] for x in preferred]
  48.         count = len(preferred)
  49.         hit = False
  50.         for c in self.connections:
  51.             u = c.get_upload()
  52.             if c in preferred:
  53.                 u.unchoke()
  54.             else:
  55.                 if count < self.min_uploads or not hit:
  56.                     u.unchoke()
  57.                     if u.is_interested():
  58.                         count += 1
  59.                         hit = True
  60.                 else:
  61.                     u.choke()
  62.  
  63.     def connection_made(self, connection, p = None):
  64.         if p is None:
  65.             p = randrange(-2, len(self.connections) + 1)
  66.         self.connections.insert(max(p, 0), connection)
  67.         self._rechoke()
  68.  
  69.     def connection_lost(self, connection):
  70.         self.connections.remove(connection)
  71.         if connection.get_upload().is_interested() and not connection.get_upload().is_choked():
  72.             self._rechoke()
  73.  
  74.     def interested(self, connection):
  75.         if not connection.get_upload().is_choked():
  76.             self._rechoke()
  77.  
  78.     def not_interested(self, connection):
  79.         if not connection.get_upload().is_choked():
  80.             self._rechoke()
  81.  
  82.     def change_max_uploads(self, newval):
  83.         def foo(self=self, newval=newval):
  84.             self._change_max_uploads(newval)
  85.         self.schedule(foo, 0);
  86.         
  87.     def _change_max_uploads(self, newval):
  88.         self.max_uploads = newval
  89.         self._rechoke()
  90.  
  91. class DummyScheduler:
  92.     def __init__(self):
  93.         self.s = []
  94.  
  95.     def __call__(self, func, delay):
  96.         self.s.append((func, delay))
  97.  
  98. class DummyConnection:
  99.     def __init__(self, v = 0):
  100.         self.u = DummyUploader()
  101.         self.d = DummyDownloader(self)
  102.         self.v = v
  103.     
  104.     def get_upload(self):
  105.         return self.u
  106.  
  107.     def get_download(self):
  108.         return self.d
  109.  
  110. class DummyDownloader:
  111.     def __init__(self, c):
  112.         self.s = False
  113.         self.c = c
  114.  
  115.     def is_snubbed(self):
  116.         return self.s
  117.  
  118.     def get_rate(self):
  119.         return self.c.v
  120.  
  121. class DummyUploader:
  122.     def __init__(self):
  123.         self.i = False
  124.         self.c = True
  125.  
  126.     def choke(self):
  127.         if not self.c:
  128.             self.c = True
  129.  
  130.     def unchoke(self):
  131.         if self.c:
  132.             self.c = False
  133.  
  134.     def is_choked(self):
  135.         return self.c
  136.  
  137.     def is_interested(self):
  138.         return self.i
  139.  
  140. def test_round_robin_with_no_downloads():
  141.     s = DummyScheduler()
  142.     Choker(2, s)
  143.     assert len(s.s) == 1
  144.     assert s.s[0][1] == 10
  145.     s.s[0][0]()
  146.     del s.s[0]
  147.     assert len(s.s) == 1
  148.     assert s.s[0][1] == 10
  149.     s.s[0][0]()
  150.     del s.s[0]
  151.     s.s[0][0]()
  152.     del s.s[0]
  153.     s.s[0][0]()
  154.     del s.s[0]
  155.  
  156. def test_resort():
  157.     s = DummyScheduler()
  158.     choker = Choker(1, s)
  159.     c1 = DummyConnection()
  160.     c2 = DummyConnection(1)
  161.     c3 = DummyConnection(2)
  162.     c4 = DummyConnection(3)
  163.     c2.u.i = True
  164.     c3.u.i = True
  165.     choker.connection_made(c1)
  166.     assert not c1.u.c
  167.     choker.connection_made(c2, 1)
  168.     assert not c1.u.c
  169.     assert not c2.u.c
  170.     choker.connection_made(c3, 1)
  171.     assert not c1.u.c
  172.     assert c2.u.c
  173.     assert not c3.u.c
  174.     c2.v = 2
  175.     c3.v = 1
  176.     choker.connection_made(c4, 1)
  177.     assert not c1.u.c
  178.     assert c2.u.c
  179.     assert not c3.u.c
  180.     assert not c4.u.c
  181.     choker.connection_lost(c4)
  182.     assert not c1.u.c
  183.     assert c2.u.c
  184.     assert not c3.u.c
  185.     s.s[0][0]()
  186.     assert not c1.u.c
  187.     assert c2.u.c
  188.     assert not c3.u.c
  189.  
  190. def test_interest():
  191.     s = DummyScheduler()
  192.     choker = Choker(1, s)
  193.     c1 = DummyConnection()
  194.     c2 = DummyConnection(1)
  195.     c3 = DummyConnection(2)
  196.     c2.u.i = True
  197.     c3.u.i = True
  198.     choker.connection_made(c1)
  199.     assert not c1.u.c
  200.     choker.connection_made(c2, 1)
  201.     assert not c1.u.c
  202.     assert not c2.u.c
  203.     choker.connection_made(c3, 1)
  204.     assert not c1.u.c
  205.     assert c2.u.c
  206.     assert not c3.u.c
  207.     c3.u.i = False
  208.     choker.not_interested(c3)
  209.     assert not c1.u.c
  210.     assert not c2.u.c
  211.     assert not c3.u.c
  212.     c3.u.i = True
  213.     choker.interested(c3)
  214.     assert not c1.u.c
  215.     assert c2.u.c
  216.     assert not c3.u.c
  217.     choker.connection_lost(c3)
  218.     assert not c1.u.c
  219.     assert not c2.u.c
  220.  
  221. def test_robin_interest():
  222.     s = DummyScheduler()
  223.     choker = Choker(1, s)
  224.     c1 = DummyConnection(0)
  225.     c2 = DummyConnection(1)
  226.     c1.u.i = True
  227.     choker.connection_made(c2)
  228.     assert not c2.u.c
  229.     choker.connection_made(c1, 0)
  230.     assert not c1.u.c
  231.     assert c2.u.c
  232.     c1.u.i = False
  233.     choker.not_interested(c1)
  234.     assert not c1.u.c
  235.     assert not c2.u.c
  236.     c1.u.i = True
  237.     choker.interested(c1)
  238.     assert not c1.u.c
  239.     assert c2.u.c
  240.     choker.connection_lost(c1)
  241.     assert not c2.u.c
  242.  
  243. def test_skip_not_interested():
  244.     s = DummyScheduler()
  245.     choker = Choker(1, s)
  246.     c1 = DummyConnection(0)
  247.     c2 = DummyConnection(1)
  248.     c3 = DummyConnection(2)
  249.     c1.u.i = True
  250.     c3.u.i = True
  251.     choker.connection_made(c2)
  252.     assert not c2.u.c
  253.     choker.connection_made(c1, 0)
  254.     assert not c1.u.c
  255.     assert c2.u.c
  256.     choker.connection_made(c3, 2)
  257.     assert not c1.u.c
  258.     assert c2.u.c
  259.     assert c3.u.c
  260.     f = s.s[0][0]
  261.     f()
  262.     assert not c1.u.c
  263.     assert c2.u.c
  264.     assert c3.u.c
  265.     f()
  266.     assert not c1.u.c
  267.     assert c2.u.c
  268.     assert c3.u.c
  269.     f()
  270.     assert c1.u.c
  271.     assert c2.u.c
  272.     assert not c3.u.c
  273.  
  274. def test_connection_lost_no_interrupt():
  275.     s = DummyScheduler()
  276.     choker = Choker(1, s)
  277.     c1 = DummyConnection(0)
  278.     c2 = DummyConnection(1)
  279.     c3 = DummyConnection(2)
  280.     c1.u.i = True
  281.     c2.u.i = True
  282.     c3.u.i = True
  283.     choker.connection_made(c1)
  284.     choker.connection_made(c2, 1)
  285.     choker.connection_made(c3, 2)
  286.     f = s.s[0][0]
  287.     f()
  288.     assert not c1.u.c
  289.     assert c2.u.c
  290.     assert c3.u.c
  291.     f()
  292.     assert not c1.u.c
  293.     assert c2.u.c
  294.     assert c3.u.c
  295.     f()
  296.     assert c1.u.c
  297.     assert not c2.u.c
  298.     assert c3.u.c
  299.     f()
  300.     assert c1.u.c
  301.     assert not c2.u.c
  302.     assert c3.u.c
  303.     f()
  304.     assert c1.u.c
  305.     assert not c2.u.c
  306.     assert c3.u.c
  307.     choker.connection_lost(c3)
  308.     assert c1.u.c
  309.     assert not c2.u.c
  310.     f()
  311.     assert not c1.u.c
  312.     assert c2.u.c
  313.     choker.connection_lost(c2)
  314.     assert not c1.u.c
  315.  
  316. def test_connection_made_no_interrupt():
  317.     s = DummyScheduler()
  318.     choker = Choker(1, s)
  319.     c1 = DummyConnection(0)
  320.     c2 = DummyConnection(1)
  321.     c3 = DummyConnection(2)
  322.     c1.u.i = True
  323.     c2.u.i = True
  324.     c3.u.i = True
  325.     choker.connection_made(c1)
  326.     choker.connection_made(c2, 1)
  327.     f = s.s[0][0]
  328.     assert not c1.u.c
  329.     assert c2.u.c
  330.     f()
  331.     assert not c1.u.c
  332.     assert c2.u.c
  333.     f()
  334.     assert not c1.u.c
  335.     assert c2.u.c
  336.     choker.connection_made(c3, 1)
  337.     assert not c1.u.c
  338.     assert c2.u.c
  339.     assert c3.u.c
  340.     f()
  341.     assert c1.u.c
  342.     assert c2.u.c
  343.     assert not c3.u.c
  344.  
  345. def test_round_robin():
  346.     s = DummyScheduler()
  347.     choker = Choker(1, s)
  348.     c1 = DummyConnection(0)
  349.     c2 = DummyConnection(1)
  350.     c1.u.i = True
  351.     c2.u.i = True
  352.     choker.connection_made(c1)
  353.     choker.connection_made(c2, 1)
  354.     f = s.s[0][0]
  355.     assert not c1.u.c
  356.     assert c2.u.c
  357.     f()
  358.     assert not c1.u.c
  359.     assert c2.u.c
  360.     f()
  361.     assert not c1.u.c
  362.     assert c2.u.c
  363.     f()
  364.     assert c1.u.c
  365.     assert not c2.u.c
  366.     f()
  367.     assert c1.u.c
  368.     assert not c2.u.c
  369.     f()
  370.     assert c1.u.c
  371.     assert not c2.u.c
  372.     f()
  373.     assert not c1.u.c
  374.     assert c2.u.c
  375.     
  376. def test_multi():
  377.     s = DummyScheduler()
  378.     choker = Choker(4, s)
  379.     c1 = DummyConnection(0)
  380.     c2 = DummyConnection(0)
  381.     c3 = DummyConnection(0)
  382.     c4 = DummyConnection(8)
  383.     c5 = DummyConnection(0)
  384.     c6 = DummyConnection(0)
  385.     c7 = DummyConnection(6)
  386.     c8 = DummyConnection(0)
  387.     c9 = DummyConnection(9)
  388.     c10 = DummyConnection(7)
  389.     c11 = DummyConnection(10)
  390.     choker.connection_made(c1, 0)
  391.     choker.connection_made(c2, 1)
  392.     choker.connection_made(c3, 2)
  393.     choker.connection_made(c4, 3)
  394.     choker.connection_made(c5, 4)
  395.     choker.connection_made(c6, 5)
  396.     choker.connection_made(c7, 6)
  397.     choker.connection_made(c8, 7)
  398.     choker.connection_made(c9, 8)
  399.     choker.connection_made(c10, 9)
  400.     choker.connection_made(c11, 10)
  401.     c2.u.i = True
  402.     c4.u.i = True
  403.     c6.u.i = True
  404.     c8.u.i = True
  405.     c10.u.i = True
  406.     c2.d.s = True
  407.     c6.d.s = True
  408.     c8.d.s = True
  409.     s.s[0][0]()
  410.     assert not c1.u.c
  411.     assert not c2.u.c
  412.     assert not c3.u.c
  413.     assert not c4.u.c
  414.     assert not c5.u.c
  415.     assert not c6.u.c
  416.     assert c7.u.c
  417.     assert c8.u.c
  418.     assert c9.u.c
  419.     assert not c10.u.c
  420.     assert c11.u.c
  421.  
  422.  
  423.